# db_util.py
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
import logging

from Util.encryption_utils import decrypt_data


class DBUtil:
    def __init__(self, db_url):
        self.engine = create_engine(db_url, echo=False)
        self.Session = sessionmaker(bind=self.engine)

    def __init__(self, environment, db_name):
        if environment == "prod":
            decrypted_db_url = decrypt_data('encrypted_data.txt', 'secret.key')
            db_url = f"{decrypted_db_url}/{db_name}"
        elif environment == "test":
            decrypted_db_url = decrypt_data('encrypted_data_test.txt', 'secret.key')
            db_url = f"{decrypted_db_url}/{db_name}"
        else:
            raise ValueError("Invalid environment specified")

        self.engine = create_engine(db_url, echo=False)
        self.Session = sessionmaker(bind=self.engine)

    def execute_query(self, query, params=None):
        session = self.Session()
        try:
            result = session.execute(query, params)
            session.commit()
            return result.fetchall()
        except Exception as e:
            session.rollback()
            logging.error(f"Query execution failed: {e}")
            return None
        finally:
            session.close()

    def execute_update(self, query, params=None):
        session = self.Session()
        try:
            session.execute(query, params)
            session.commit()
        except Exception as e:
            session.rollback()
            logging.error(f"Update execution failed: {e}")
        finally:
            session.close()
